"""
Peak finding script for HEMC validation.
Identifies peak positions in energy deposition profiles.
"""

import csv
import glob
from pathlib import Path
from typing import cast

import matplotlib.pyplot as plt
import numpy as np
from scipy.signal import find_peaks

from pyauto.config.router import router

from ..config.models import HEMCConfig

_MODULENAME = Path(__file__).stem
_EDEP_DUMP_PATTERN = "eDep_*.npz"


def read_energy_deposition_data(npz_file: Path) -> tuple:
    """
    Read energy deposition data from npz file.

    Args:
        npz_file: Path to CSV file

    Returns:
        Tuple of (iZ array, mean energy deposition array)
    """
    iz, mean_edep = [], []
    with open(npz_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            iz_, _, _, sumv, _, n = line.split(",")
            iz.append(int(iz_))
            aver = float(sumv) / int(n) if int(n) > 0 else 0.0
            mean_edep.append(aver)
    return np.array(iz), np.array(mean_edep)


def find_peaks_in_profile(
    mean_edep: np.ndarray, min_height_ratio: float = 0.05, min_distance: int = 3
) -> np.ndarray:
    """
    Find peaks in energy deposition profile.

    Args:
        mean_edep: Mean energy deposition array
        min_height_ratio: Minimum peak height as ratio of (max - min)
        min_distance: Minimum distance between peaks

    Returns:
        Array of peak indices
    """
    h_min, h_max = mean_edep.min(), mean_edep.max()
    min_height = h_min + min_height_ratio * (h_max - h_min)
    peaks, _ = find_peaks(mean_edep, height=min_height, distance=min_distance)
    return peaks


def plot_peaks_with_annotations(
    iz: np.ndarray,
    mean_edep: np.ndarray,
    peaks: np.ndarray,
    save_path: Path = None,
    title: str = "Peak positions & values",
) -> None:
    """
    Plot energy deposition profile with peak annotations.

    Args:
        iz: iZ position array
        mean_edep: Mean energy deposition array
        peaks: Peak indices
        save_path: Path to save figure (None to not save)
        title: Title for the plot
    """
    ## 为了画面整洁, 只画出最高的峰
    ## 找出 mean_dep 中最大值的位置（在 peaks 中的索引）
    max_idx_in_peaks = np.argmax(mean_edep[peaks])
    ## 对应的 peaks 值
    peak_corresponding_to_max = peaks[max_idx_in_peaks]
    peaks_sd = [peak_corresponding_to_max]

    plt.figure(figsize=(8, 4))
    plt.plot(iz, mean_edep, lw=1.2, label="Mean energy deposit")
    plt.scatter(iz[peaks_sd], mean_edep[peaks_sd], color="r", zorder=5, label="Peaks")

    ##=========================== 加竖线 Add vertical lines at peak positions
    for p in iz[peaks_sd]:
        plt.axvline(p, ls="--", color="gray", lw=0.8)

    #
    # Annotate peaks with values
    for p in peaks_sd:
        plt.annotate(
            f"{iz[p]}; {mean_edep[p]:.0f}",
            xy=(iz[p], mean_edep[p]),
            xytext=(5, 10),
            textcoords="offset points",
            color="red",
            fontsize=14,
            alpha=0.6,
            ha="left",
            va="bottom",
            arrowprops=dict(arrowstyle="->", color="red", lw=0.8),
        )

    plt.xlabel("iZ")
    plt.ylabel("Energy deposit / MeV")
    plt.title(title)
    plt.tight_layout()
    plt.legend()

    if save_path:
        plt.savefig(save_path, dpi=200)
        print(f"Saved annotated figure -> {save_path}")

    # plt.show()


@router.register()
def findpeaks(config: HEMCConfig):
    """Main entry point for peak finding."""
    ## 检查
    outdir = cast(Path, config.project_paths.g4outdir)
    resultdir = cast(Path, config.project_paths.resultdir)

    ## Find all eDep npz files in output directory
    pattern = str(outdir / _EDEP_DUMP_PATTERN)
    npz_files = glob.glob(pattern)

    if not npz_files:
        print(f"Error: No eDep npz files found at {outdir / _EDEP_DUMP_PATTERN}")
        return

    for npz_file_path in npz_files:
        npz_file = Path(npz_file_path)

        # Read data
        iz, mean_edep = read_energy_deposition_data(npz_file)

        # Find peaks
        peaks = find_peaks_in_profile(mean_edep)

        # Print results to console
        print(f"File: {npz_file.name}")
        print("Peak positions (iZ):", iz[peaks])
        print("Peak values  (MeV):", mean_edep[peaks])
        print("-" * 50)

        # Extract energy from filename for the title
        filename = npz_file.stem
        energy = filename.split("_")[1]  # Extract energy part (e.g., "700.0MeV")

        # Plot and save
        save_fig = resultdir / f"eDep_{energy}_peaks_ann.png"
        plot_peaks_with_annotations(
            iz,
            mean_edep,
            peaks,
            save_path=save_fig,
            title=f"eDep_{energy} - peak positions & values; mm-MeV",
        )

        ##000000000000000000000000000000000000 Save peaks data to CSV
        save_csv = resultdir / f"eDep_{energy}_peaks.csv"
        with open(save_csv, "w", newline="", encoding="utf-8") as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(["# iZ", "Energy_deposit_MeV"])  # Header
            for peak_idx in peaks:
                writer.writerow([iz[peak_idx], mean_edep[peak_idx]])
        print(f"Saved peaks data -> {save_csv}")


if __name__ == "__main__":
    pass
